import abc
from dataclasses import dataclass
from typing import Iterable

import numpy as np
import torch
from transformers import Conv1D, PreTrainedModel

from zkl_rome import ComputeVDeltaCallback, ComputeVDeltaHparams, ComputeVDeltaMetrics, TokenizedPreserving, \
    TokenizedRewriting, compute_left_right


@dataclass(kw_only=True)
class ComputeWilkeHparams:
    acc_threshold: float | None = None
    loss_threshold: float | None = None
    w_delta_norm_ratio_threshold: float | None = None


@dataclass(kw_only=True)
class ComputeWilkeMetrics(ComputeVDeltaMetrics):
    processed_steps_num: int
    rewriting_acc: torch.Tensor
    rewriting_loss: torch.Tensor
    preserving_loss: torch.Tensor
    regularization_loss: torch.Tensor
    loss: torch.Tensor
    w_delta_norm_ratio: torch.Tensor


class ComputeWilkeCallback(abc.ABC):
    @abc.abstractmethod
    def on_start(self, modules_num: int, hparams: ComputeWilkeHparams):
        pass

    @abc.abstractmethod
    def on_traverse(self, module_index: int, metrics: ComputeWilkeMetrics) -> bool | None:
        pass

    @abc.abstractmethod
    def on_stop(self):
        pass


def compute_wilke(*,
    model: PreTrainedModel,
    modules: Iterable[torch.nn.Module],
    rewriting: TokenizedRewriting,
    prefixes: Iterable[np.ndarray],
    preservings: Iterable[TokenizedPreserving],
    modules_c_inv: Iterable[torch.Tensor] | None = None,
    compute_v_delta_hparams: ComputeVDeltaHparams,
    compute_v_delta_callback: ComputeVDeltaCallback | None = None,
    hparams: ComputeWilkeHparams | None = None,
    callback: ComputeWilkeCallback | None = None,
) -> tuple[int, torch.Tensor, torch.Tensor]:
    modules = tuple(modules)
    modules_c_inv = tuple(modules_c_inv) if modules_c_inv is not None else ([None] * len(modules))

    if hparams is None:
        hparams = ComputeWilkeHparams()

    if callback is not None:
        callback.on_start(len(modules), hparams)

    results = []
    for index, (module, c_inv) in enumerate(zip(modules, modules_c_inv)):
        with torch.no_grad():
            w = extract_w(module)
            w_norm = torch.norm(w)

        compute_v_delta_callback = WrappedComputeVDeltaCallback(compute_v_delta_callback)

        (left, right) = compute_left_right(
            model=model,
            module=module,
            rewriting=rewriting,
            prefixes=prefixes,
            preservings=preservings,
            c_inv=c_inv,
            compute_v_delta_hparams=compute_v_delta_hparams,
            compute_v_delta_callback=compute_v_delta_callback)

        v_delta_metrics = compute_v_delta_callback.last_metrics
        assert isinstance(v_delta_metrics, ComputeVDeltaMetrics)

        with torch.no_grad():
            w_delta = torch.outer(left, right)
            w_delta_norm = torch.norm(w_delta)
            w_delta_norm_ratio = w_delta_norm / w_norm

        if callback is not None:
            metrics = ComputeWilkeMetrics(
                processed_steps_num=v_delta_metrics.processed_steps_num,
                rewriting_acc=v_delta_metrics.rewriting_acc,
                rewriting_loss=v_delta_metrics.rewriting_loss,
                preserving_loss=v_delta_metrics.preserving_loss,
                regularization_loss=v_delta_metrics.regularization_loss,
                loss=v_delta_metrics.loss,
                w_delta_norm_ratio=w_delta_norm_ratio)
            callback_result = callback.on_traverse(index, metrics)
            if callback_result is not None and not callback_result:
                continue

        if hparams is not None:
            if hparams.acc_threshold is not None:
                if v_delta_metrics.rewriting_acc.mean().item() < hparams.acc_threshold:
                    continue
            if hparams.loss_threshold is not None:
                if v_delta_metrics.loss.item() > hparams.loss_threshold:
                    continue
            if hparams.w_delta_norm_ratio_threshold is not None:
                if w_delta_norm_ratio.item() > hparams.w_delta_norm_ratio_threshold:
                    continue

        results.append((index, left, right, w_delta_norm_ratio))

    if callback is not None:
        callback.on_stop()

    if not results:
        raise RuntimeError("No candidates satisfied the metrics filter!")

    best_index, best_left, best_right, _ = \
        min(results, key=lambda x: x[-1])
    return best_index, best_left, best_right


def extract_w(module: torch.nn.Module) -> torch.Tensor:
    if isinstance(module, Conv1D):
        return module.weight
    elif isinstance(module, (torch.nn.Linear, torch.nn.LazyLinear)):
        return module.weight.T
    else:
        raise TypeError(f"Unrecognized module {module}. Unable to get its original weight!")


class WrappedComputeVDeltaCallback(ComputeVDeltaCallback):
    def __init__(self, wrapped: ComputeVDeltaCallback | None):
        self.wrapped = wrapped
        self.last_metrics: ComputeVDeltaMetrics | None = None

    def on_start(self, hparams: ComputeVDeltaHparams):
        if self.wrapped is not None:
            self.wrapped.on_start(hparams)

    def on_step(self, metrics: ComputeVDeltaMetrics):
        if self.wrapped is not None:
            self.wrapped.on_step(metrics)

    def on_stop(self, metrics: ComputeVDeltaMetrics):
        if self.wrapped is not None:
            self.wrapped.on_stop(metrics)
        self.last_metrics = metrics
